import csv
import random

# Define consolidated vote_core_configs for each year
# This dictionary stores the configuration for each year's vote generation, including:
# - party_weights: The desired distribution of votes across parties.
# - year_favorites: The favored policies for specific parties in that year.
# - vote1 to vote5: The policy options for each of the 5 votes a voter can cast in a year.
vote_core_configs = {
    2025: {
        'party_weights': {1: 800, 2: 4600, 3: 900, 4: 1700, 5: 2000},
        'year_favorites': {2: [7, 9]},
        'vote1': {1: 1, 2: 6, 3: 11, 4: 16, 5: 21},
        'vote2': {1: 2, 2: 7, 3: 12, 4: 17, 5: 22},
        'vote3': {1: 3, 2: 8, 3: 13, 4: 18, 5: 23},
        'vote4': {1: 4, 2: 9, 3: 14, 4: 19, 5: 24},
        'vote5': {1: 5, 2: 10, 3: 15, 4: 20, 5: 25}
    },
    2026: {
        'party_weights': {1: 1100, 2: 4400, 3: 1700, 4: 1300, 5: 1500},
        'year_favorites': {2: [32, 33, 34]},
        'vote1': {1: 26, 2: 31, 3: 36, 4: 41, 5: 46},
        'vote2': {1: 27, 2: 32, 3: 37, 4: 42, 5: 47},
        'vote3': {1: 28, 2: 33, 3: 38, 4: 43, 5: 48},
        'vote4': {1: 29, 2: 34, 3: 39, 4: 44, 5: 49},
        'vote5': {1: 30, 2: 35, 3: 40, 4: 45, 5: 50}
    },
    2027: {
        'party_weights': {1: 600, 2: 3300, 3: 1400, 4: 2600, 5: 2100},
        'year_favorites': {2: [57, 58, 59], 4: [67, 68, 69]},
        'vote1': {1: 51, 2: 56, 3: 61, 4: 66, 5: 71},
        'vote2': {1: 52, 2: 57, 3: 62, 4: 67, 5: 72},
        'vote3': {1: 53, 2: 58, 3: 63, 4: 68, 5: 73},
        'vote4': {1: 54, 2: 59, 3: 64, 4: 69, 5: 74},
        'vote5': {1: 55, 2: 60, 3: 65, 4: 70, 5: 75}
    },
    2028: {
        'party_weights': {1: 700, 2: 3100, 3: 2500, 4: 2200, 5: 1500},
        'year_favorites': {2: [82, 83, 84], 3: [87, 88, 89], 4: [92, 93, 94]},
        'vote1': {1: 76, 2: 81, 3: 86, 4: 91, 5: 96},
        'vote2': {1: 77, 2: 82, 3: 87, 4: 92, 5: 97},
        'vote3': {1: 78, 2: 83, 3: 88, 4: 93, 5: 98},
        'vote4': {1: 79, 2: 84, 3: 89, 4: 94, 5: 99},
        'vote5': {1: 80, 2: 85, 3: 90, 4: 95, 5: 100}
    },
    2029: {
        'party_weights': {1: 800, 2: 2400, 3: 2000, 4: 3100, 5: 1700},
        'year_favorites': {2: [107, 108], 3: [112, 113], 4: [117, 118, 119]},
        'vote1': {1: 101, 2: 106, 3: 111, 4: 116, 5: 121},
        'vote2': {1: 102, 2: 107, 3: 112, 4: 117, 5: 122},
        'vote3': {1: 103, 2: 108, 3: 113, 4: 118, 5: 123},
        'vote4': {1: 104, 2: 109, 3: 114, 4: 119, 5: 124},
        'vote5': {1: 105, 2: 110, 3: 115, 4: 120, 5: 125}
    }
}


def generate_votes():
    """
    Generates a list of votes based on the vote_core_configs.

    Returns:
        list: A list of votes, where each vote is a list containing:
            - VoteID (unique identifier)
            - VoterID
            - VoteCycleYear
            - PolicyID
            - PartyID
    """
    votes = []
    vote_id = 1

    for year, year_data in vote_core_configs.items():

        party_weights = year_data['party_weights']
        year_favorites = year_data['year_favorites']

        # Initialize vote allocations from party_weights
        party_allocations = party_weights.copy()
        favorite_allocations = {party: int(party_weights[party] * 0.52) for party in year_favorites}
        fva = sum(favorite_allocations.values())    # Total number of favorite votes to allocate

        for voter_id in range(1, 2001):  # Loop through all 2k voters

            for vote_num in range(1, 6):  # Ensure all 5 votes are cast

                vote_key = 'vote%d' % vote_num
                policy_options = year_data[vote_key]    # Policy options for current year
                available_favorites = []

                if fva > 0:

                    for partyid, policy_ids in year_favorites.items():
                        if favorite_allocations[partyid] > 0:
                            # Check if any favorite policies are available in this vote
                            available_favorites = [
                                policy_id
                                for policy_id in policy_ids
                                if policy_id in policy_options.values()
                            ]
                            if available_favorites:
                                party_id = partyid
                                policy_id = policy_options[party_id]

                                favorite_allocations[party_id] -= 1 # Removes the favorite vote from the party
                                fva -= 1    # Decrement the total favorite allocation count
                                break   # Exit the inner loop once a favorite is allocated

                # If no favorites were allocated or they're exhausted, distribute remaining votes                
                if fva == 0 or not available_favorites:

                    # Finds parties with available vote slots to register a vote
                    available_parties = [
                        party_id
                        for party_id, allocation in party_allocations.items()
                        if allocation > 0
                    ]
                    if available_parties:
                        party_id = random.choice(available_parties)
                        policy_id = policy_options[party_id]

                # Registers the vote and removes a party vote from the ballot
                party_allocations[party_id] -= 1
                votes.append([vote_id, voter_id, year, policy_id, party_id])
                vote_id += 1


    return votes


def write_csv(votes, filename):
    """
    Writes the generated votes to a CSV file.
    Args:
        votes (list): The list of votes to write.
        filename (str): The name of the CSV file to create.
    """
    with open(filename, 'wb') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['VoteID', 'VoterID', 'VoteCycleYear', 'PolicyID', 'PartyID'])
        writer.writerows(votes)


def print_vote_distribution(votes):
    """
    Prints the vote distribution per party for each year.
    Args:
        votes (list): The list of votes to analyze.
    """
    vote_counts = {}
    for _, _, year, _, party_id in votes:
        if year not in vote_counts:
            vote_counts[year] = {1: 0, 2: 0, 3: 0, 4: 0, 5: 0}
        vote_counts[year][party_id] += 1

    for year, counts in vote_counts.iteritems():  # Use iteritems() for Python 2.7
        print "Vote distribution for year %d:" % year 
        for party, count in counts.iteritems():
            print "  Party %d: %d votes" % (party, count)
        print ""

# Initializer
if __name__ == '__main__':
    votes = generate_votes()
    write_csv(votes, 'rawlvotes.csv')
    print "CSV file 'rawlvotes.csv' has been created with %d votes." % len(votes)
    print ""
    print_vote_distribution(votes)    
